#pragma once

#include "RungeKutta.h"
#include "MultigridSolver.h"

/**
 * @file AdvectionDiffusion.h
 * @author xingyifan
 * @date 2022-08-31 22:52
 *
 * @description: 对流扩散方程求解器
 */

template <int Dim>
class AdvectionDiffusionSolver
{
public:
    // 注意使用初始化列表时要注意不同值初始化的先后顺序问题，这里使用 -k * nu * 0.25 而不是 -k * nu * gamma
    AdvectionDiffusionSolver(FuncXt f, Real k, Real nu, GhostFiller<Dim> &ghostfiller)
        : t_(0), f_(f), k_(k), nu_(nu), msolver_("I-Laplacian", k * nu * 0.25, ghostfiller) {}

    /**
     * @file AdvectionDiffusion.h
     * @author xingyifan
     * @date 2022-08-31 22:52
     *
     * @description: 传入初值，速度和推进次数，返回运行时间
     */

    Real solve(Grid<Dim> &U, GridX<Dim> &V, int N)
    {
        // 开始计时
        clock_t start = clock();

        // 循环计算每一步，移动 N 次
        for (int i = 0; i < N; i++)
        {
            shift(U, V);
        }

        // 结束计时
        Real t = Real(clock() - start) / CLOCKS_PER_SEC;

        // 将时间归零
        t_ = 0;

        return t;
    }

    /**
     * @file AdvectionDiffusion.h
     * @author xingyifan
     * @date 2022-08-31 22:53
     *
     * @description: 使用理查德外推计算更精确解
     */

    Real exact(Grid<Dim> &U, GridX<Dim> &V, int N)
    {
        // 开始计时
        clock_t start = clock();

        // 循环计算每一步，移动 N 次
        for (int i = 0; i < N; i++)
        {
            extrapolation(U, V);
        }

        // 结束计时
        Real t = Real(clock() - start) / CLOCKS_PER_SEC;

        // 将时间归零
        t_ = 0;

        return t;
    }

private:
    /**
     * @file AdvectionDiffusion.h
     * @author xingyifan
     * @date 2022-08-31 22:53
     *
     * @description: phi_t = -D(u phi) + nu L(phi) + f
     */

    void shift(Grid<Dim> &U, GridX<Dim> &V)
    {
        // 存放算子结果
        std::vector<Grid<Dim>> Div(6);
        std::vector<Grid<Dim>> Lap(6);

        // 压力项
        std::vector<Grid<Dim>> f(6, Grid<Dim>(U.width()));

        // 创建算子
        Laplacian<Dim> laplacian;
        Divergence<Dim> divergence;
        Integrator<3, Dim> integrator;

        // 创建单位区域
        Domain<Dim> domain;

        // 填充器
        GhostFiller<Dim> &ghostfiller = msolver_.ghostfiller_;

        // RK 方法的系数
        Matrix &AI = imexrungekutta_.AI_;
        Matrix &AE = imexrungekutta_.AE_;
        Vector &b = imexrungekutta_.b_;
        Vector &c = imexrungekutta_.c_;

        // 临时容器
        Grid<Dim> Phi = U;

        // 循环求解 (I-k nu gamma L)X = F
        for (int i = 1; i < 6; i++)
        {
            // 记录 F
            Grid<Dim> F = U;
            for (int j = 0; j < i - 1; j++)
            {
                F += TensorOp::add(Div[j] * (AE[i * 6 + j] * k_), Lap[j] * (AI[i * 6 + j] * nu_ * k_));
            }

            // 包装函数
            FuncX F_ = [=](Vector x)
            {
                return f_(x, t_ + c[i - 1] * k_);
            };
            // 填充
            integrator.fill(domain, f[i - 1], F_);

            // 分别计算 f-D<u>, L<u>
            Div[i - 1] = std::move(TensorOp::sub(f[i - 1], divergence(Phi, V, ghostfiller, t_ + c[i - 1] * k_)));
            Lap[i - 1] = std::move(laplacian(Phi, ghostfiller, t_ + c[i - 1] * k_));

            F += TensorOp::add(Div[i - 1] * (AE[i * 6 + i - 1] * k_), Lap[i - 1] * (AI[i * 6 + i - 1] * nu_ * k_));

            // 修正 f，补充边界条件
            laplacian.supplement(F, ghostfiller, t_ + c[i] * k_, msolver_.coeff_);

            // 使用多重网格求解
            msolver_.solve(F, Phi);
        }

        // 包装函数
        FuncX F_ = [=](Vector x)
        {
            return f_(x, t_ + c[5] * k_);
        };
        // 填充
        integrator.fill(domain, f[5], F_);

        // 计算最后一项
        Div[5] = std::move(TensorOp::sub(f[5], divergence(Phi, V, ghostfiller, t_ + c[5] * k_)));
        Lap[5] = std::move(laplacian(Phi, ghostfiller, t_ + c[5] * k_));

        // 加权求和
        for (int i = 0; i < 6; i++)
        {
            U += TensorOp::add(Div[i] * (b[i] * k_), Lap[i] * (b[i] * nu_ * k_));
        }

        // 时间推进
        t_ += k_;
    }

    /**
     * @file AdvectionDiffusion.h
     * @author xingyifan
     * @date 2022-08-31 22:57
     *
     * @description: 理查德外推过程
     */

    void extrapolation(Grid<Dim> &U, GridX<Dim> &V)
    {
        // 记录初值
        Grid<Dim> tU = U;

        // 计算一步
        shift(U, V);

        // 回退时间并改变步长
        t_ -= k_;
        k_ /= 2;

        // 注意 coeff = k nu gamma 也要减半
        msolver_.coeff_ /= 2;

        // 计算两步
        shift(tU, V);
        shift(tU, V);

        // 恢复步长
        k_ *= 2;
        msolver_.coeff_ *= 2;

        // 使用更精确解，精度 p = 3
        U = TensorOp::add(tU, TensorOp::divide(tU - U, pow(2, 3) - 1));
    }

private:
    FuncXt f_; // 压力项

    Real t_;  // 当前时间
    Real k_;  // 时间步长
    Real nu_; // 扩散常数

    IMEXRungeKutta imexrungekutta_; // 隐式-显式迭代
    MultigridSolver<Dim> msolver_;  // 多重网格求解器
};